#!/usr/bin/env python3
"""
稳定版无人机导航系统模拟 V4.2 - PID 输出为期望加速度 + 传感器低通滤波
说明:
 - PIDController.compute 现在返回期望加速度 (m/s^2)，并按轴限幅
 - Drone.update 使用 total_force = mass * (a_cmd - gravity)
 - 在 GPS 测量上加入指数平滑，减少测量跳变
 - 初始若干步 (起飞保护) 限制最大加速度，避免仿真初期剧烈动作
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.gridspec as gridspec
import math
import time
from queue import PriorityQueue
import cvxpy as cp
import warnings

# 忽略特定警告
warnings.filterwarnings("ignore", category=UserWarning, module="cvxpy")
warnings.filterwarnings("ignore", category=RuntimeWarning)

# 设置matplotlib使用支持中文的字体（如系统有 SimHei）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 调试日志级别
DEBUG_LEVEL = 1  # 0: 无调试, 1: 基本调试, 2: 详细调试, 3: 非常详细


def debug_log(level, message):
    """调试日志函数"""
    if DEBUG_LEVEL >= level:
        print(f"[DEBUG{level}] {message}")


class PIDController:
    """PID 控制器 — 现在输出期望加速度 (m/s^2)"""

    def __init__(self, kp, ki, kd, max_output=float('inf')):
        # kp, ki, kd 可以是长度为3的数组（对每轴分别设置）
        self.kp = np.array(kp, dtype=float)
        self.ki = np.array(ki, dtype=float)
        self.kd = np.array(kd, dtype=float)
        self.max_output = max_output  # 单轴最大加速度限制（m/s^2）

        self.integral = np.zeros(3)
        self.previous_error = np.zeros(3)
        self.previous_time = None  # 改为 None，更好处理第一次调用
        self.error_history = []

    def compute(self, setpoint, current_value, accel_max=5.0):
        """计算期望加速度（m/s^2），并限制幅值。
           setpoint, current_value: 都是 3D 向量（位置）
           accel_max: 单轴最大加速度（m/s^2）
        """
        current_time = time.time()
        # 第一次调用或时间步过小则用一个小 dt 避免除零
        if self.previous_time is None:
            dt = 0.01
        else:
            dt = current_time - self.previous_time
            if dt <= 1e-6:
                dt = 0.01

        error = np.array(setpoint, dtype=float) - np.array(current_value, dtype=float)
        self.error_history.append(np.linalg.norm(error))

        # 比例项（注意 kp 的尺度使得 kp*error 给出加速度量级）
        p_term = self.kp * error

        # 积分项
        self.integral += error * dt
        i_term = self.ki * self.integral

        # 微分项
        derivative = (error - self.previous_error) / dt
        d_term = self.kd * derivative

        # 期望加速度
        a_cmd = p_term + i_term + d_term

        # 限幅（按轴）
        a_cmd = np.clip(a_cmd, -accel_max, accel_max)

        # 更新状态
        self.previous_error = error
        self.previous_time = current_time

        debug_log(3, f"PID 期望加速度: {a_cmd}, 误差: {np.linalg.norm(error):.4f}, dt={dt:.4f}")

        return a_cmd

    def get_error_history(self):
        return self.error_history


class SimpleMPCController:
    """简化版模型预测控制器（保留原结构）——注: 返回的 control_output 我们当成期望加速度 a_cmd 使用"""
    def __init__(self, mass, horizon=5, dt=0.1):
        self.mass = mass
        self.horizon = horizon
        self.dt = dt

        # 简单系统矩阵（这里保持原始简单模型）
        self.A = np.eye(3)
        self.B = np.eye(3) * dt

        # 权重矩阵（可调整）
        self.Q = np.diag([10.0, 10.0, 15.0])
        self.R = np.diag([0.1, 0.1, 0.1])

        # 控制输入约束（如果把 u 当作加速度，则范围设置适中）
        self.u_min = -3.0
        self.u_max = 3.0

        self.x_min = np.array([-np.inf, -np.inf, 0.1])
        self.x_max = np.array([np.inf, np.inf, np.inf])

        # 优化变量
        self.x = cp.Variable((3, horizon + 1))
        self.u = cp.Variable((3, horizon))

        self.reference = None
        self.error_history = []

        # 备用 PID（当 MPC 失败时回退）
        self.backup_pid = PIDController(kp=np.array([4.0, 4.0, 6.0]),
                                        ki=np.array([0.05, 0.05, 0.1]),
                                        kd=np.array([2.0, 2.0, 3.0]),
                                        max_output=20.0)

        self.solve_count = 0
        self.fail_count = 0

    def set_reference(self, reference):
        """设置参考轨迹，reference 期望形状 (3, horizon+1)"""
        if reference.shape[0] != 3 or reference.shape[1] != self.horizon + 1:
            debug_log(1, f"参考轨迹维度错误: {reference.shape}，期望: (3, {self.horizon + 1})")
            # 尝试修正
            if reference.shape[0] != 3:
                reference = reference[:3, :]
            if reference.shape[1] < self.horizon + 1:
                last_point = reference[:, -1].reshape(3, 1)
                padding = np.repeat(last_point, self.horizon + 1 - reference.shape[1], axis=1)
                reference = np.hstack([reference, padding])
            else:
                reference = reference[:, :self.horizon + 1]

        self.reference = reference
        debug_log(2, f"设置参考轨迹: {self.reference[:,0]}")

    def compute(self, current_state):
        """计算控制输出（这里把输出视作期望加速度 a_cmd）"""
        # 校验输入
        if np.any(np.isnan(current_state)) or np.any(np.isinf(current_state)):
            debug_log(1, "当前状态包含 NaN/Inf，使用备用 PID")
            return self.backup_pid.compute(self.reference[:, 0] if self.reference is not None else np.zeros(3),
                                           current_state[0:3] if len(current_state) > 3 else current_state)

        if self.reference is None:
            debug_log(1, "参考轨迹未设置，返回零加速度")
            return np.zeros(3)

        # x0 仅为位置（因为我们在此保持 3 维状态）
        x0 = current_state[0:3] if len(current_state) >= 3 else current_state

        # 构建成本与约束（简化）
        cost = 0
        constraints = [self.x[:, 0] == x0]

        for t in range(self.horizon):
            cost += cp.quad_form(self.x[:, t] - self.reference[:, t], self.Q)
            cost += cp.quad_form(self.u[:, t], self.R)
            constraints += [self.x[:, t + 1] == self.A @ self.x[:, t] + self.B @ self.u[:, t]]
            constraints += [self.u[:, t] >= self.u_min, self.u[:, t] <= self.u_max]
            constraints += [self.x[:, t] >= self.x_min, self.x[:, t] <= self.x_max]

        cost += cp.quad_form(self.x[:, self.horizon] - self.reference[:, self.horizon], self.Q * 2)
        problem = cp.Problem(cp.Minimize(cost), constraints)

        self.solve_count += 1

        try:
            problem.solve(solver=cp.ECOS, verbose=False, max_iters=100, abstol=1e-3, reltol=1e-3)
        except Exception as e:
            try:
                problem.solve(solver=cp.OSQP, verbose=False, max_iter=4000)
            except Exception as e2:
                self.fail_count += 1
                debug_log(1, f"MPC 求解器全部失败: {e}, {e2}")
                return self.backup_pid.compute(self.reference[:, 0], x0)

        if problem.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
            self.fail_count += 1
            debug_log(1, f"MPC 求解失败! 状态: {problem.status}")
            return self.backup_pid.compute(self.reference[:, 0], x0)

        # 记录误差
        error = np.linalg.norm(x0 - self.reference[:, 0])
        self.error_history.append(error)

        control_output = self.u[:, 0].value
        if control_output is None or not np.all(np.isfinite(control_output)):
            self.fail_count += 1
            debug_log(1, f"MPC 输出无效: {control_output}")
            return self.backup_pid.compute(self.reference[:, 0], x0)

        debug_log(2, f"MPC 控制输出 (treated as a_cmd): {control_output}, 误差: {error:.4f}")
        return np.array(control_output, dtype=float)

    def get_error_history(self):
        return self.error_history


class SensorNoise:
    """传感器噪声模型（GPS/IMU）"""

    def __init__(self):
        # GPS 噪声参数 (位置测量) - 适度缩小以避免初始跳变
        self.gps_position_std = 0.05  # 米
        self.gps_velocity_std = 0.03  # 米/秒

        # IMU 噪声参数 (加速度测量)
        self.accel_bias = np.array([0.01, -0.02, 0.005])
        self.accel_std = 0.02

        # 陀螺仪噪声参数
        self.gyro_bias = np.array([0.005, -0.003, 0.008])
        self.gyro_std = 0.01

    def add_gps_noise(self, position, velocity):
        noisy_position = position + np.random.normal(0, self.gps_position_std, 3)
        noisy_velocity = velocity + np.random.normal(0, self.gps_velocity_std, 3)
        return noisy_position, noisy_velocity

    def add_imu_noise(self, acceleration):
        noisy_acceleration = acceleration + self.accel_bias + np.random.normal(0, self.accel_std, 3)
        return noisy_acceleration

    def add_gyro_noise(self, angular_velocity):
        noisy_angular_velocity = angular_velocity + self.gyro_bias + np.random.normal(0, self.gyro_std, 3)
        return noisy_angular_velocity


class AStarPlanner:
    """A*路径规划（保持你原实现）"""

    def __init__(self, grid_size=1.0):
        self.grid_size = grid_size

    def heuristic(self, a, b):
        return abs(a[0] - b[0]) + abs(a[1] - b[1]) + abs(a[2] - b[2])

    def plan(self, start, goal, obstacles):
        start_grid = (int(start[0] / self.grid_size),
                      int(start[1] / self.grid_size),
                      int(start[2] / self.grid_size))
        goal_grid = (int(goal[0] / self.grid_size),
                     int(goal[1] / self.grid_size),
                     int(goal[2] / self.grid_size))

        obstacle_grid = set()
        for obs in obstacles:
            obs_grid = (int(obs[0] / self.grid_size),
                        int(obs[1] / self.grid_size),
                        int(obs[2] / self.grid_size))
            obstacle_grid.add(obs_grid)

        frontier = PriorityQueue()
        frontier.put((0, start_grid))
        came_from = {}
        cost_so_far = {}
        came_from[start_grid] = None
        cost_so_far[start_grid] = 0

        directions = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    directions.append((dx, dy, dz))

        while not frontier.empty():
            current = frontier.get()[1]
            if current == goal_grid:
                break
            for dx, dy, dz in directions:
                next_node = (current[0] + dx, current[1] + dy, current[2] + dz)
                if next_node in obstacle_grid:
                    continue
                new_cost = cost_so_far[current] + math.sqrt(dx * dx + dy * dy + dz * dz)
                if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
                    cost_so_far[next_node] = new_cost
                    priority = new_cost + self.heuristic(goal_grid, next_node)
                    frontier.put((priority, next_node))
                    came_from[next_node] = current

        path = []
        current = goal_grid
        while current != start_grid:
            path.append(current)
            if current not in came_from:
                return []
            current = came_from[current]
        path.append(start_grid)
        path.reverse()

        continuous_path = []
        for p in path:
            continuous_path.append(np.array([
                p[0] * self.grid_size + self.grid_size / 2,
                p[1] * self.grid_size + self.grid_size / 2,
                p[2] * self.grid_size + self.grid_size / 2
            ]))

        if len(continuous_path) > 2:
            smoothed_path = self.smooth_path(continuous_path, obstacles)
            return smoothed_path
        return continuous_path

    def smooth_path(self, path, obstacles):
        smoothed = np.array(path)
        alpha = 0.1
        beta = 0.3
        for _ in range(50):
            for i in range(1, len(path) - 1):
                for obs in obstacles:
                    obs_dist = np.linalg.norm(smoothed[i] - obs)
                    if obs_dist < 2.0:
                        smoothed[i] += 0.1 * (smoothed[i] - obs) / (obs_dist + 1e-6)
                smoothed[i] += alpha * (path[i] - smoothed[i]) + beta * (smoothed[i - 1] + smoothed[i + 1] - 2 * smoothed[i])
        return smoothed.tolist()


class Drone:
    def __init__(self, x=0, y=0, z=0, mass=1.0, controller_type="PID"):
        # 真实状态
        self.true_position = np.array([x, y, z], dtype=float)
        self.true_velocity = np.array([0., 0., 0.], dtype=float)
        self.true_acceleration = np.array([0., 0., 0.], dtype=float)

        # 传感器测量值（初始化为真实值）
        self.measured_position = self.true_position.copy()
        self.measured_velocity = self.true_velocity.copy()
        self.measured_acceleration = self.true_acceleration.copy()

        # 物理参数
        self.mass = mass
        self.max_thrust = 25.0  # N
        self.gravity = np.array([0., 0., -9.81])  # m/s^2

        # 控制器
        self.pid_controller = PIDController(
            kp=np.array([4.0, 4.0, 6.0]),
            ki=np.array([0.05, 0.05, 0.1]),
            kd=np.array([2.0, 2.0, 3.0]),
            max_output=20.0
        )
        self.mpc_controller = SimpleMPCController(mass=mass, horizon=5, dt=0.1)
        self.controller_type = controller_type

        # 传感器噪声模型
        self.sensor_noise = SensorNoise()

        # 轨迹记录
        self.true_trajectory = []
        self.measured_trajectory = []

        # MPC 参考轨迹（如果使用 MPC）
        self.reference_trajectory = None

        # 计数器
        self.step_count = 0

        # 低通滤波系数（GPS）
        self.gps_alpha = 0.6  # 越接近1平滑越慢，越接近0更依赖新测量

    def set_controller(self, controller_type):
        self.controller_type = controller_type
        debug_log(1, f"设置控制器: {controller_type}")

    def set_reference_trajectory(self, reference):
        """为 MPC 设置参考轨迹（按原逻辑保留）"""
        if self.controller_type == "MPC":
            waypoints_np = [np.array(wp) for wp in reference]
            references = []
            for i in range(len(waypoints_np) - 1):
                start = waypoints_np[i]
                end = waypoints_np[i + 1]
                num_points = 5
                for j in range(num_points):
                    t = j / (num_points - 1)
                    point = start + t * (end - start)
                    references.append(point)
            debug_log(2, f"Generated {len(references)} reference points for MPC")
            horizon = self.mpc_controller.horizon
            ref_states = np.zeros((3, horizon + 1))
            for i in range(horizon + 1):
                idx = min(i, len(references) - 1)
                ref_states[:, i] = references[idx]
            self.reference_trajectory = ref_states
            self.mpc_controller.set_reference(ref_states)

    def update_sensors(self):
        """更新传感器测量（添加噪声 + 低通滤波）"""
        noisy_pos, noisy_vel = self.sensor_noise.add_gps_noise(self.true_position, self.true_velocity)
        # 指数平滑（低通滤波），减少单步突变
        alpha = self.gps_alpha
        self.measured_position = alpha * self.measured_position + (1 - alpha) * noisy_pos
        self.measured_velocity = alpha * self.measured_velocity + (1 - alpha) * noisy_vel

        # IMU 加速度测量
        self.measured_acceleration = self.sensor_noise.add_imu_noise(self.true_acceleration)

        debug_log(3, f"传感器 - 真实位置: {self.true_position}, 原噪声位置: {noisy_pos}, 平滑后位置: {self.measured_position}")

    def update(self, target_position, dt=0.1):
        """更新无人机状态（将加速度命令转为推力）"""
        self.step_count += 1

        # 更新传感器
        self.update_sensors()

        # 起飞保护：仿真初期限制最大加速度，避免瞬态过大
        startup_steps = 8
        if self.step_count <= startup_steps:
            accel_limit = 1.0  # m/s^2 in startup
        else:
            accel_limit = 5.0  # normal

        # 计算期望加速度 a_cmd（m/s^2）
        if self.controller_type == "PID":
            a_cmd = self.pid_controller.compute(target_position, self.measured_position, accel_max=accel_limit)
            debug_log(2, f"使用 PID 控制器, 目标位置: {target_position}, a_cmd: {a_cmd}")
        else:
            # MPC 分支：我们让 MPC 返回的 u 被视为 a_cmd（如上所述）
            if self.reference_trajectory is not None:
                horizon = self.mpc_controller.horizon
                new_reference = np.zeros((3, horizon + 1))
                # 第一个点为当前位置
                new_reference[:, 0] = self.measured_position
                # 中间线性插值
                for i in range(1, horizon):
                    t = i / horizon
                    new_reference[:, i] = self.measured_position + t * (target_position - self.measured_position)
                new_reference[:, horizon] = target_position
                self.mpc_controller.set_reference(new_reference)

            current_state = np.concatenate([self.measured_position, self.measured_velocity])
            a_cmd = self.mpc_controller.compute(current_state)
            # 安全限幅（以防 MPC 输出过大）
            a_cmd = np.clip(a_cmd, -accel_limit, accel_limit)
            debug_log(2, f"使用 MPC 控制器, 目标位置: {target_position}, a_cmd: {a_cmd}")

        # 将期望加速度转换为需要的推力：F = m * (a_cmd - gravity)
        total_force = self.mass * (a_cmd - self.gravity)  # gravity 矢量为 [0,0,-9.81]
        # 限制推力
        thrust_norm = np.linalg.norm(total_force)
        if thrust_norm > 1e-6 and thrust_norm > self.max_thrust:
            total_force = total_force / thrust_norm * self.max_thrust

        # 计算真实加速度并更新物理状态
        self.true_acceleration = total_force / self.mass
        self.true_velocity += self.true_acceleration * dt
        self.true_position += self.true_velocity * dt
        self.true_position[2] = max(0.1, self.true_position[2])  # 不低于地面

        # 记录
        self.true_trajectory.append(self.true_position.copy())
        self.measured_trajectory.append(self.measured_position.copy())

        # 航点到达判断（更稳健）
        horizontal_distance = np.linalg.norm(target_position[0:2] - self.true_position[0:2])
        vertical_distance = abs(target_position[2] - self.true_position[2])
        reached = horizontal_distance < 0.5 and vertical_distance < 0.3

        # 关键 debug 输出（每 N 步）
        if self.step_count % 5 == 0:
            debug_log(1, f"步数 {self.step_count}: true_pos={self.true_position}, measured_pos={self.measured_position}, a_cmd={a_cmd}, thrust_norm={np.linalg.norm(total_force):.3f}")

        return reached

    def get_true_position(self):
        return self.true_position.copy()

    def get_measured_position(self):
        return self.measured_position.copy()

    def get_true_trajectory(self):
        return np.array(self.true_trajectory)

    def get_measured_trajectory(self):
        return np.array(self.measured_trajectory)

    def get_controller_error_history(self):
        if self.controller_type == "PID":
            return self.pid_controller.get_error_history()
        else:
            return self.mpc_controller.get_error_history()


class NavigationSystem:
    def __init__(self, controller_type="PID"):
        # 障碍物列表
        self.obstacles = [
            np.array([5, 5, 2]),
            np.array([8, 3, 3]),
            np.array([12, 8, 4]),
            np.array([15, 12, 5]),
            np.array([7, 10, 3]),
            np.array([10, 15, 4]),
            np.array([13, 5, 3]),
            np.array([18, 8, 4])
        ]

        self.start = np.array([0, 0, 1])
        self.goal = np.array([20, 15, 6])

        self.planner = AStarPlanner(grid_size=1.0)
        self.waypoints = self.plan_path()
        self.current_waypoint_idx = 1

        self.drone = Drone(x=self.start[0], y=self.start[1], z=self.start[2], mass=1.0, controller_type=controller_type)

        if controller_type == "MPC":
            self.setup_mpc_reference()

        self.mission_complete = False
        self.update_count = 0

        debug_log(1, "规划路径:")
        for i, wp in enumerate(self.waypoints):
            debug_log(1, f"航点 {i}: {wp}")

    def plan_path(self):
        debug_log(1, "Planning path...")
        path = self.planner.plan(self.start, self.goal, self.obstacles)
        if not path:
            debug_log(1, "No path found! 使用直线")
            return [self.start, self.goal]
        debug_log(1, f"Path planning completed, total {len(path)} waypoints")
        return path

    def setup_mpc_reference(self):
        waypoints_np = [np.array(wp) for wp in self.waypoints]
        references = []
        for i in range(len(waypoints_np) - 1):
            start = waypoints_np[i]
            end = waypoints_np[i + 1]
            num_points = 5
            for j in range(num_points):
                t = j / (num_points - 1)
                point = start + t * (end - start)
                references.append(point)
        debug_log(2, f"Generated {len(references)} reference points for MPC")
        self.drone.set_reference_trajectory(references)

    def update(self, dt=0.1):
        self.update_count += 1
        if self.mission_complete:
            return True

        target = self.waypoints[self.current_waypoint_idx]
        debug_log(1, f"更新 {self.update_count}: 当前目标航点 {self.current_waypoint_idx}/{len(self.waypoints) - 1}: {target}")

        reached = self.drone.update(target, dt)

        current_pos = self.drone.get_true_position()
        debug_log(1, f"当前位置: {current_pos}, 规划路径点: {target}")

        if reached:
            debug_log(1, f"到达航点 {self.current_waypoint_idx}: {target}")
            self.current_waypoint_idx += 1
            if self.current_waypoint_idx >= len(self.waypoints):
                debug_log(1, "任务完成!")
                self.mission_complete = True
                return True

        return False

    def get_drone_true_position(self):
        return self.drone.get_true_position()

    def get_drone_measured_position(self):
        return self.drone.get_measured_position()

    def get_waypoints(self):
        return self.waypoints

    def get_true_trajectory(self):
        return self.drone.get_true_trajectory()

    def get_measured_trajectory(self):
        return self.drone.get_measured_trajectory()

    def get_obstacles(self):
        return self.obstacles

    def get_controller_error_history(self):
        return self.drone.get_controller_error_history()


class Simulation:
    def __init__(self, controller_type="PID"):
        self.nav_system = NavigationSystem(controller_type)
        self.fig = plt.figure(figsize=(16, 12))
        gs = gridspec.GridSpec(3, 2, figure=self.fig)

        self.ax_3d = self.fig.add_subplot(gs[0:2, 0], projection='3d')
        self.ax_error = self.fig.add_subplot(gs[0, 1])
        self.ax_control = self.fig.add_subplot(gs[1, 1])
        self.ax_sensor = self.fig.add_subplot(gs[2, 0])
        self.ax_performance = self.fig.add_subplot(gs[2, 1])

        self.ax_3d.set_title(f'Drone Navigation Simulation - {controller_type} Controller')
        self.ax_3d.set_xlabel('X (m)')
        self.ax_3d.set_ylabel('Y (m)')
        self.ax_3d.set_zlabel('Z (m)')
        self.ax_3d.set_xlim(-2, 25)
        self.ax_3d.set_ylim(-2, 20)
        self.ax_3d.set_zlim(0, 10)

        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax_3d.scatter(start[0], start[1], start[2], c='green', marker='o', s=100, label='Start')
        self.ax_3d.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100, label='Goal')

        waypoints = self.nav_system.get_waypoints()
        wp_x = [wp[0] for wp in waypoints]
        wp_y = [wp[1] for wp in waypoints]
        wp_z = [wp[2] for wp in waypoints]
        self.ax_3d.plot(wp_x, wp_y, wp_z, 'r--', alpha=0.5, label='Planned Path')

        obstacles = self.nav_system.get_obstacles()
        for obs in obstacles:
            self.draw_cube(obs, 1.0, 'gray')

        self.true_marker, = self.ax_3d.plot([], [], [], 'bo', markersize=8, label='True Position')
        self.measured_marker, = self.ax_3d.plot([], [], [], 'rx', markersize=6, label='Measured Position')

        self.true_trajectory_line, = self.ax_3d.plot([], [], [], 'b-', alpha=0.7, linewidth=2, label='True Trajectory')
        self.measured_trajectory_line, = self.ax_3d.plot([], [], [], 'r-', alpha=0.5, linewidth=1, label='Measured Trajectory')

        self.ax_error.set_title('Controller Error')
        self.ax_error.set_xlabel('Time Step')
        self.ax_error.set_ylabel('Error (m)')
        self.ax_error.grid(True)
        self.error_line, = self.ax_error.plot([], [], 'b-')

        self.ax_control.set_title('Control Input')
        self.ax_control.set_xlabel('Time Step')
        self.ax_control.set_ylabel('Control Force (N)')
        self.ax_control.grid(True)
        self.control_lines = [
            self.ax_control.plot([], [], label='Fx')[0],
            self.ax_control.plot([], [], label='Fy')[0],
            self.ax_control.plot([], [], label='Fz')[0]
        ]

        self.ax_sensor.set_title('Sensor Measurement vs True Value')
        self.ax_sensor.set_xlabel('Time Step')
        self.ax_sensor.set_ylabel('Position (m)')
        self.ax_sensor.grid(True)
        self.sensor_lines = [
            self.ax_sensor.plot([], [], label='True X')[0],
            self.ax_sensor.plot([], [], label='Measured X')[0],
            self.ax_sensor.plot([], [], label='True Y')[0],
            self.ax_sensor.plot([], [], label='Measured Y')[0]
        ]

        self.ax_performance.set_title('Performance Metrics')
        self.ax_performance.set_xlabel('Metric')
        self.ax_performance.set_ylabel('Value')
        self.ax_performance.grid(True)

        self.ax_3d.legend()
        self.ax_control.legend()
        self.ax_sensor.legend()
        self.ax_3d.view_init(elev=20, azim=45)

        self.animation = None

        self.error_history = []
        self.control_history = []
        self.true_position_history = []
        self.measured_position_history = []
        self.time_steps = []

    def draw_cube(self, center, size, color):
        r = size / 2
        x = [center[0] - r, center[0] + r]
        y = [center[1] - r, center[1] + r]
        z = [center[2] - r, center[2] + r]
        from itertools import product
        vertices = list(product(x, y, z))
        edges = [
            (0, 1), (0, 2), (0, 4),
            (1, 3), (1, 5),
            (2, 3), (2, 6),
            (3, 7),
            (4, 5), (4, 6),
            (5, 7),
            (6, 7)
        ]
        for edge in edges:
            start = vertices[edge[0]]
            end = vertices[edge[1]]
            self.ax_3d.plot3D(*zip(start, end), color=color, alpha=0.6)

    def update_plot(self, frame):
        complete = self.nav_system.update(dt=0.1)

        true_pos = self.nav_system.get_drone_true_position()
        measured_pos = self.nav_system.get_drone_measured_position()

        self.true_marker.set_data([true_pos[0]], [true_pos[1]])
        self.true_marker.set_3d_properties([true_pos[2]])

        self.measured_marker.set_data([measured_pos[0]], [measured_pos[1]])
        self.measured_marker.set_3d_properties([measured_pos[2]])

        true_trajectory = self.nav_system.get_true_trajectory()
        measured_trajectory = self.nav_system.get_measured_trajectory()

        if len(true_trajectory) > 1:
            self.true_trajectory_line.set_data(true_trajectory[:, 0], true_trajectory[:, 1])
            self.true_trajectory_line.set_3d_properties(true_trajectory[:, 2])

        if len(measured_trajectory) > 1:
            self.measured_trajectory_line.set_data(measured_trajectory[:, 0], measured_trajectory[:, 1])
            self.measured_trajectory_line.set_3d_properties(measured_trajectory[:, 2])

        error_history = self.nav_system.get_controller_error_history()
        if error_history:
            self.error_line.set_data(range(len(error_history)), error_history)
            self.ax_error.relim()
            self.ax_error.autoscale_view()

        self.time_steps.append(len(self.time_steps))
        self.true_position_history.append(true_pos)
        self.measured_position_history.append(measured_pos)

        recent_points = min(50, len(self.time_steps))
        if recent_points > 0:
            time_window = self.time_steps[-recent_points:]
            true_x = [p[0] for p in self.true_position_history[-recent_points:]]
            true_y = [p[1] for p in self.true_position_history[-recent_points:]]
            measured_x = [p[0] for p in self.measured_position_history[-recent_points:]]
            measured_y = [p[1] for p in self.measured_position_history[-recent_points:]]
            self.sensor_lines[0].set_data(time_window, true_x)
            self.sensor_lines[1].set_data(time_window, measured_x)
            self.sensor_lines[2].set_data(time_window, true_y)
            self.sensor_lines[3].set_data(time_window, measured_y)
            self.ax_sensor.relim()
            self.ax_sensor.autoscale_view()

        if complete:
            self.animation.event_source.stop()
            debug_log(1, "Simulation completed")
            self.calculate_performance()

        return (self.true_marker, self.measured_marker,
                self.true_trajectory_line, self.measured_trajectory_line,
                self.error_line, *self.sensor_lines)

    def calculate_performance(self):
        true_trajectory = self.nav_system.get_true_trajectory()
        waypoints = self.nav_system.get_waypoints()
        if len(true_trajectory) < 2:
            return
        path_length = 0
        for i in range(1, len(true_trajectory)):
            path_length += np.linalg.norm(true_trajectory[i] - true_trajectory[i - 1])
        straight_distance = np.linalg.norm(waypoints[-1] - waypoints[0])
        efficiency = straight_distance / path_length if path_length > 0 else 0
        error_history = self.nav_system.get_controller_error_history()
        avg_error = np.mean(error_history) if error_history else 0
        max_error = np.max(error_history) if error_history else 0
        metrics = ['Path Length', 'Straight Distance', 'Efficiency', 'Avg Error', 'Max Error']
        values = [path_length, straight_distance, efficiency, avg_error, max_error]
        self.ax_performance.bar(metrics, values)
        self.ax_performance.tick_params(axis='x', rotation=45)
        for i, v in enumerate(values):
            self.ax_performance.text(i, v + 0.01, f'{v:.2f}', ha='center')
        debug_log(1, f"性能指标:")
        debug_log(1, f"  路径长度: {path_length:.2f} m")
        debug_log(1, f"  直线距离: {straight_distance:.2f} m")
        debug_log(1, f"  效率: {efficiency:.2f}")
        debug_log(1, f"  平均误差: {avg_error:.2f} m")
        debug_log(1, f"  最大误差: {max_error:.2f} m")

    def run(self):
        debug_log(1, "Starting drone navigation simulation...")
        debug_log(1, f"Controller type: {self.nav_system.drone.controller_type}")
        debug_log(1, f"Start point: {self.nav_system.start}")
        debug_log(1, f"Goal point: {self.nav_system.goal}")
        debug_log(1, f"Number of obstacles: {len(self.nav_system.get_obstacles())}")

        self.animation = FuncAnimation(
            self.fig, self.update_plot, frames=500,
            interval=50, blit=True, repeat=False
        )
        plt.tight_layout()
        plt.show()


# 直接运行两个控制器以便比较
if __name__ == "__main__":
    debug_log(1, "First testing with PID controller to ensure basic functionality...")
    sim_pid = Simulation(controller_type="PID")
    sim_pid.run()

    debug_log(1, "Now testing with MPC controller...")
    sim_mpc = Simulation(controller_type="MPC")
    sim_mpc.run()
